import cv2
from PIL import Image, ImageDraw, ImageFont, ImageFile
import numpy as np

# 允许Pillow加载被截断的图像文件
ImageFile.LOAD_TRUNCATED_IMAGES = True
from ultralytics import YOLO
from sklearn.linear_model import RANSACRegressor
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline
import torch
import warnings
from sklearn.exceptions import UndefinedMetricWarning
import logging
from concurrent.futures import ThreadPoolExecutor

logger = logging.getLogger(__name__)

# 全局变量，用于存储加载的字体，避免重复加载
_global_font = None
_global_layer_font = None


def get_yolo_model(model_path):
    """
    加载YOLOv8模型，并自动选择GPU或CPU。
    同时加载字体文件，避免重复加载。
    """
    global _global_font, _global_layer_font

    # 检查是否有可用的CUDA设备
    device = "cuda" if torch.cuda.is_available() else "cpu"
    logger.info(f"正在使用设备: {device} 加载模型: {model_path}")
    model = YOLO(model_path)
    model.to(device)

    # 优化：全局加载字体，避免重复I/O和内存分配
    if _global_font is None:
        try:
            _global_font = ImageFont.truetype(
                "static/fonts/SourceHanSansSC-Regular.otf", 12
            )
            _global_layer_font = ImageFont.truetype(
                "static/fonts/SourceHanSansSC-Regular.otf", 20
            )
            logger.info("字体文件已成功加载。")
        except IOError:
            logger.warning("中文字体加载失败，将回退到默认字体。")
            _global_font = ImageFont.load_default()
            _global_layer_font = ImageFont.load_default()

    return model


def detect_layers(
    boxes,
    image_width,
    min_bottles_per_layer=3,
    residual_threshold=25,
    slope_tolerance=0.1,
    vertical_stack_threshold=50,
):
    """
    使用基于Y坐标聚类、RANSAC和全局斜率一致性约束的鲁棒分层检测算法。
    新增了对垂直堆叠瓶子的处理逻辑。
    :param boxes: YOLO检测结果的boxes对象。
    :param min_bottles_per_layer: 形成一层所需的最少瓶子数。
    :param residual_threshold: Y坐标聚类的垂直距离阈值。
    :param slope_tolerance: 全局斜率一致性检查的容差。
    :param vertical_stack_threshold: 判断是否为垂直堆叠的x坐标重叠阈值。
    :return: 一个包含每层信息的列表，确保所有层大致平行且不相交。
    """
    if len(boxes) < min_bottles_per_layer:
        return []

    # --- 1. 提取所有瓶子的底部中心点，并动态计算聚类阈值 ---
    boxes_cpu = boxes.xyxy.cpu()

    # 根据瓶子高度动态计算Y坐标聚类阈值
    heights = boxes_cpu[:, 3] - boxes_cpu[:, 1]
    if len(heights) > 0:
        # 使用中位数高度来避免异常值（如检测不全的瓶子）的影响
        median_height = np.median(heights)
        # 根据用户反馈，当瓶子底部坐标偏差小于高度的1/4时，仍视为一层
        y_cluster_threshold = median_height / 4.0
    else:
        # 如果没有检测到瓶子，则使用默认值
        y_cluster_threshold = residual_threshold

    x_centers = (boxes_cpu[:, 0] + boxes_cpu[:, 2]) / 2
    y_bottoms = boxes_cpu[:, 3]
    indices = np.arange(len(boxes))
    points = np.stack((x_centers, y_bottoms, indices), axis=1)

    # --- 2. 基于Y坐标对瓶子进行聚类 (Y-coalescing) ---
    sorted_points = points[points[:, 1].argsort()]
    clusters = []
    if len(sorted_points) > 0:
        current_cluster = [sorted_points[0]]
        for i in range(1, len(sorted_points)):
            if sorted_points[i][1] - current_cluster[-1][1] < y_cluster_threshold:
                current_cluster.append(sorted_points[i])
            else:
                if len(current_cluster) >= min_bottles_per_layer:
                    clusters.append(np.array(current_cluster))
                current_cluster = [sorted_points[i]]
        if len(current_cluster) >= min_bottles_per_layer:
            clusters.append(np.array(current_cluster))

    # --- 3. 对每个聚类运行RANSAC，找出潜在的层 ---
    potential_layers = []
    for cluster in clusters:
        X = cluster[:, 0].reshape(-1, 1)
        y = cluster[:, 1]
        if len(X) < 2:
            continue

        ransac = RANSACRegressor(
            residual_threshold=residual_threshold / 2,  # RANSAC使用更严格的阈值
            min_samples=max(2, int(min_bottles_per_layer * 0.5)),
            random_state=42,
            max_trials=50,  # 优化：限制最大尝试次数
        )
        try:
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=UndefinedMetricWarning)
                ransac.fit(X, y)

            if hasattr(ransac.estimator_, "coef_"):
                inlier_mask = ransac.inlier_mask_
                # 确保内点数量也满足最小要求
                if np.sum(inlier_mask) >= min_bottles_per_layer:
                    potential_layers.append(
                        {
                            "model": ransac,
                            "bottle_indices": cluster[inlier_mask, 2]
                            .astype(int)
                            .tolist(),
                            "avg_y": np.mean(cluster[inlier_mask, 1]),
                            "slope": ransac.estimator_.coef_[0],
                        }
                    )
        except ValueError:
            logger.warning("RANSAC在聚类上拟合失败，已跳过。")
            continue

    if not potential_layers:
        return []

    # --- 4. 全局斜率一致性约束 (找到所有水平层) ---
    slopes = [layer["slope"] for layer in potential_layers]
    if not slopes:
        return []
    median_slope = np.median(slopes)

    horizontal_layers = []
    for layer in potential_layers:
        if abs(layer["slope"] - median_slope) < slope_tolerance:
            horizontal_layers.append(layer)

    if not horizontal_layers:
        return []

    # --- 4.5. 后处理：过滤叠放层 (性能优化版) ---
    # 通过向量化操作替换多重循环，显著提升性能
    if not horizontal_layers:
        final_layers = []
    else:
        all_boxes_np = boxes.xyxy.cpu().numpy()
        all_bottom_centers_x = (all_boxes_np[:, 0] + all_boxes_np[:, 2]) / 2
        all_bottom_y = all_boxes_np[:, 3]

        # 按Y坐标均值排序，这样只需检查当前层是否位于其下方的任何层之上
        sorted_layers = sorted(horizontal_layers, key=lambda l: l["avg_y"])
        is_stacked = [False] * len(sorted_layers)

        for i, layer_i in enumerate(sorted_layers):
            # 如果一个层已经被确定为堆叠层，则跳过
            if is_stacked[i]:
                continue

            # 获取当前层i的瓶子底部中心点
            indices_i = np.array(layer_i["bottle_indices"])
            bottom_centers_x_i = all_bottom_centers_x[indices_i]
            bottom_y_i = all_bottom_y[indices_i]

            # 仅与位于其下方的层进行比较
            for j in range(i + 1, len(sorted_layers)):
                layer_j = sorted_layers[j]
                indices_j = np.array(layer_j["bottle_indices"])
                boxes_j = all_boxes_np[indices_j]

                # 向量化检查：判断i层的任何一个瓶子底部是否在j层的任何一个瓶子框内
                # 使用广播机制 (None) 来创建(len(i), len(j))的比较矩阵
                x_overlap = (bottom_centers_x_i[:, None] > boxes_j[:, 0]) & (
                    bottom_centers_x_i[:, None] < boxes_j[:, 2]
                )
                y_overlap = (bottom_y_i[:, None] > boxes_j[:, 1]) & (
                    bottom_y_i[:, None] < boxes_j[:, 3]
                )

                # 如果存在任何重叠，则认为i层是堆叠层
                if (x_overlap & y_overlap).any():
                    is_stacked[i] = True
                    # 找到一个堆叠关系后，即可停止对j的搜索，继续检查下一个i
                    break

        # 仅保留未被标记为堆叠的层
        final_layers = [
            layer for idx, layer in enumerate(sorted_layers) if not is_stacked[idx]
        ]

    # --- 4.6. 后处理：过滤相交的层，只保留更水平的线 ---
    if len(final_layers) > 1:
        layers_to_keep = [True] * len(final_layers)
        for i in range(len(final_layers)):
            for j in range(i + 1, len(final_layers)):
                if not layers_to_keep[i] or not layers_to_keep[j]:
                    continue

                layer_i = final_layers[i]
                layer_j = final_layers[j]

                m_i = layer_i["slope"]
                c_i = layer_i["model"].predict([[0]])[0]
                m_j = layer_j["slope"]
                c_j = layer_j["model"].predict([[0]])[0]

                # 避免除以零
                if abs(m_i - m_j) < 1e-6:
                    continue

                # 计算交点x坐标
                x_intersect = (c_j - c_i) / (m_i - m_j)

                # 如果交点在图像范围内，则移除更不水平的线
                if 0 <= x_intersect <= image_width:
                    if abs(m_i) > abs(m_j):
                        layers_to_keep[i] = False
                    else:
                        layers_to_keep[j] = False

        final_layers = [
            final_layers[i] for i in range(len(final_layers)) if layers_to_keep[i]
        ]

    # --- 5. 按垂直位置排序并返回 ---
    final_layers.sort(key=lambda l: l["avg_y"])
    return final_layers


def run_detection(model, image, conf, iou):
    """在单独的线程中运行模型推理，并自动启用半精度以提升性能。"""
    # 如果有可用的CUDA设备，则自动启用FP16半精度推理，可以显著提升性能
    half = torch.cuda.is_available()
    return model(image, conf=conf, iou=iou, verbose=False, half=half)[0]


def process_all_detections(
    bottle_model,
    sku_model,
    image_path,
    output_path,
    bottle_conf=0.25,
    bottle_iou=0.45,
    sku_conf=0.25,
    sku_iou=0.45,
    slope_tolerance=0.1,
    clarity_threshold=50,
):
    """
    在一个函数中处理所有检测和绘制任务，并优化性能。
    1. 使用OpenCV高效读取图像。
    2. 并行进行瓶子和SKU检测，自动启用半精度。
    3. 优化绘制流程，减少数据转换。
    4. 保存最终图像并返回所有统计数据。
    """
    # --- 1. 使用OpenCV高效读取图像并检查清晰度 ---
    image_cv = cv2.imread(image_path)
    if image_cv is None:
        logger.error(f"无法打开或读取图像文件: {image_path}")
        return None, "无法读取图像文件。"

    clarity = check_image_clarity(image_cv)
    if clarity < clarity_threshold:
        logger.warning(
            f"图像 {image_path} 清晰度过低 ({clarity:.2f})，低于阈值 {clarity_threshold}"
        )
        return None, "图片太模糊，请重新上传或拍摄。"

    image_for_bottle_sku_detection = image_cv

    # --- 2. 并行推理 ---
    with ThreadPoolExecutor(max_workers=2) as executor:
        # 将OpenCV图像（BGR格式）直接传递给推理函数
        future_bottle = executor.submit(
            run_detection,
            bottle_model,
            image_for_bottle_sku_detection,
            bottle_conf,
            bottle_iou,
        )
        future_sku = executor.submit(
            run_detection, sku_model, image_for_bottle_sku_detection, sku_conf, sku_iou
        )
        bottle_results = future_bottle.result()
        sku_results = future_sku.result()

    # --- 3. 准备绘制 ---
    # 使用已全局加载的字体
    font = _global_font
    layer_font = _global_layer_font
    if font is None or layer_font is None:
        logger.error("字体未加载，无法进行绘制。")
        return None, "内部错误：字体未加载。"

    # --- 3. 准备绘制 ---
    # 将OpenCV图像转换为PIL图像以便绘制文本和复杂图形
    image_pil = Image.fromarray(cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(image_pil)
    image_width = image_pil.width

    # 使用已全局加载的字体
    font = _global_font
    layer_font = _global_layer_font
    if font is None or layer_font is None:
        logger.error("字体未加载，无法进行绘制。")
        return None, "内部错误：字体未加载。"

    offset_x, offset_y = (0, 0)  # 不再有冰柜，偏移量始终为0

    # --- 5. 绘制瓶子中心点 (使用PIL) ---
    # 注意：这里的bottle_results.boxes是相对于image_for_bottle_sku_detection的
    for box in bottle_results.boxes:
        x1, y1, x2, y2 = box.xyxy[0].cpu()
        # 不再有冰柜偏移，直接使用原始坐标
        x1_orig, y1_orig, x2_orig, y2_orig = x1, y1, x2, y2
        center_x = int((x1_orig + x2_orig) / 2)
        center_y = int((y1_orig + y2_orig) / 2)
        # 在PIL图像上绘制绿色点
        draw.ellipse(
            (center_x - 5, center_y - 5, center_x + 5, center_y + 5),
            fill=(0, 255, 0),
            outline=(0, 255, 0),
        )

    # --- 6. 绘制SKU边界框和分层信息 (需要PIL) ---
    # 绘制SKU
    sku_names = sku_results.names
    colors = {
        i: tuple(np.random.randint(0, 255, 3).tolist()) for i in range(len(sku_names))
    }
    sku_counts = {}
    for box in sku_results.boxes:
        x1, y1, x2, y2 = [int(i) for i in box.xyxy[0].cpu()]
        # 不再有冰柜偏移，直接使用原始坐标
        x1_orig, y1_orig, x2_orig, y2_orig = x1, y1, x2, y2

        cls_id = int(box.cls[0].cpu())
        class_name = sku_names[cls_id]
        sku_counts[class_name] = sku_counts.get(class_name, 0) + 1
        color = colors.get(cls_id, (255, 0, 0))  # 默认红色
        draw.rectangle([x1_orig, y1_orig, x2_orig, y2_orig], outline=color, width=2)
        label = f"{class_name}"
        text_bbox = draw.textbbox((0, 0), label, font=font)
        text_w, text_h = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
        draw.rectangle(
            [x1_orig, y1_orig - text_h - 5, x1_orig + text_w, y1_orig - 5], fill=color
        )
        draw.text(
            (x1_orig, y1_orig - text_h - 5), label, fill=(255, 255, 255), font=font
        )

    # 绘制分层
    # 不再有冰柜，直接使用原始的bottle_results.boxes
    layers = detect_layers(
        bottle_results.boxes,
        image_width,
        slope_tolerance=slope_tolerance,
        vertical_stack_threshold=50,
    )

    for i, layer in enumerate(layers):
        model = layer["model"]
        y0 = model.predict([[0]])[0]
        y1 = model.predict([[image_width]])[0]
        draw.line(
            [(0, int(y0)), (image_width, int(y1))], fill=(255, 0, 0), width=3
        )  # 红色线
        label = f"第 {i + 1} 层"
        text_bbox = draw.textbbox((0, 0), label, font=layer_font)
        text_w = text_bbox[2] - text_bbox[0]
        draw.text(
            (image_width - text_w - 10, int(y1) - 30),
            label,
            fill=(255, 0, 0),
            font=layer_font,
        )

    # --- 7. 保存和返回 ---
    final_image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_path, final_image_cv)

    # 统计瓶子和SKU数量，只统计冰柜内部的
    sku_count = len(sku_results.boxes)
    layer_count = len(layers)

    # --- 8. 按层统计SKU (更正逻辑) ---
    layers_details = []
    if layers:
        image_height = image_pil.height

        # 定义所有分界线的y坐标，包括图像顶部和底部
        y_boundaries = [0]  # 图像顶部
        for layer in layers:
            y_boundary = int(layer["model"].predict([[image_width / 2]])[0])
            y_boundary = max(0, min(y_boundary, image_height - 1))
            y_boundaries.append(y_boundary)
        y_boundaries.append(image_height)  # 图像底部

        # 确保边界唯一并排序，以防万一
        y_boundaries = sorted(list(set(y_boundaries)))

        # 遍历每个由分界线定义的层区间
        for i in range(len(y_boundaries) - 1):
            y_start = y_boundaries[i]
            y_end = y_boundaries[i + 1]
            layer_index = i + 1  # 层号从1开始

            layer_skus_total = 0
            layer_skus_details = {}

            # 统计该层的SKU
            for box in sku_results.boxes:
                # 不再有冰柜偏移，直接使用原始坐标进行判断
                y_center_orig = int((box.xyxy[0][1] + box.xyxy[0][3]) / 2)
                if y_start <= y_center_orig < y_end:
                    layer_skus_total += 1
                    class_name = sku_names[int(box.cls[0])]
                    layer_skus_details[class_name] = (
                        layer_skus_details.get(class_name, 0) + 1
                    )

            # 如果该层有内容，则添加到结果列表
            if layer_skus_total > 0:
                layers_details.append(
                    {
                        "layer_index": layer_index,
                        "sku_total": layer_skus_total,
                        "sku_details": layer_skus_details,
                    }
                )

    # 确保总瓶子数和SKU数是基于裁剪区域的
    # 如果没有冰柜，则这些值是基于全图的
    # 如果有冰柜，这些值是基于冰柜内部的
    # 这里的bottle_count和sku_count已经是在裁剪区域内的数量，无需额外调整

    return (
        sku_count,
        sku_counts,
        layer_count,
        layers_details,
    )


def check_image_clarity(image_cv):
    """
    检查图像清晰度 (接收OpenCV图像对象)
    """
    gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)
    clarity = cv2.Laplacian(gray, cv2.CV_64F).var()
    return clarity
